{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "former-supply",
   "metadata": {},
   "source": [
    "# Using Flux.jl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "guilty-underground",
   "metadata": {},
   "outputs": [],
   "source": [
    "using Flux\r\n",
    "using MLDatasets: Cora"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "continued-circuit",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = Cora.dataset();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "minimal-reply",
   "metadata": {},
   "outputs": [],
   "source": [
    "using SparseArrays\r\n",
    "using LinearAlgebra:I"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "covered-saturn",
   "metadata": {},
   "outputs": [],
   "source": [
    "include(\"../src/utils.jl\");\n",
    "include(\"../src/layers/base.jl\");\n",
    "include(\"../src/layers/conv.jl\");\n",
    "include(\"../src/models/conv.jl\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "strange-reviewer",
   "metadata": {},
   "outputs": [],
   "source": [
    "struct GATConv <: GraphLayer\n",
    "    linear\n",
    "    attenᵤ\n",
    "    attenᵥ\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "dependent-blend",
   "metadata": {},
   "outputs": [],
   "source": [
    "GATConv(in_feats::Integer, out_feats::Integer;σ=identity,bias=true) = \n",
    "    GATConv(Dense(in_feats,out_feats,σ;bias=bias),Dense(out_feats,1;bias=false),Dense(out_feats,1;bias=false))\n",
    "\n",
    "function (m::GATConv)(X,A)\n",
    "    H = m.linear(X)\n",
    "    αᵤ = m.attenᵤ(H)\n",
    "    αᵥ = m.attenᵥ(H)\n",
    "    α = αᵤ .+  αᵥ'\n",
    "    α[A .== 0] .= -9e15\n",
    "    α = softmax(leakyrelu.(α))\n",
    "    return H*α\n",
    "end\n",
    "\n",
    "Flux.@functor GATConv\n",
    "Flux.params(m::GATConv) = params(m.linear,m.attenᵤ,m.attenᵥ)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "polished-study",
   "metadata": {},
   "outputs": [
    {
     "ename": "LoadError",
     "evalue": "Mutating arrays is not supported -- called copyto!(::SubArray{Float64, 1, SparseMatrixCSC{Float64, Int64}, Tuple{Vector{CartesianIndex{2}}}, false}, _...)",
     "output_type": "error",
     "traceback": [
      "Mutating arrays is not supported -- called copyto!(::SubArray{Float64, 1, SparseMatrixCSC{Float64, Int64}, Tuple{Vector{CartesianIndex{2}}}, false}, _...)",
      "",
      "Stacktrace:",
      "  [1] error(s::String)",
      "    @ Base .\\error.jl:33",
      "  [2] (::Zygote.var\"#436#437\"{SubArray{Float64, 1, SparseMatrixCSC{Float64, Int64}, Tuple{Vector{CartesianIndex{2}}}, false}})(#unused#::Nothing)",
      "    @ Zygote ~\\.julia\\packages\\Zygote\\rv6db\\src\\lib\\array.jl:74",
      "  [3] (::Zygote.var\"#2331#back#438\"{Zygote.var\"#436#437\"{SubArray{Float64, 1, SparseMatrixCSC{Float64, Int64}, Tuple{Vector{CartesianIndex{2}}}, false}}})(Δ::Nothing)",
      "    @ Zygote ~\\.julia\\packages\\ZygoteRules\\AIbCs\\src\\adjoint.jl:67",
      "  [4] Pullback",
      "    @ .\\broadcast.jl:894 [inlined]",
      "  [5] Pullback",
      "    @ .\\broadcast.jl:891 [inlined]",
      "  [6] Pullback",
      "    @ .\\broadcast.jl:887 [inlined]",
      "  [7] Pullback",
      "    @ .\\In[12]:9 [inlined]",
      "  [8] (::typeof(∂(λ)))(Δ::Matrix{Float64})",
      "    @ Zygote ~\\.julia\\packages\\Zygote\\rv6db\\src\\compiler\\interface2.jl:0",
      "  [9] Pullback",
      "    @ ~\\Documents\\CodeSpace\\SimpleGNN\\src\\models\\conv.jl:19 [inlined]",
      " [10] (::typeof(∂(λ)))(Δ::Matrix{Float64})",
      "    @ Zygote ~\\.julia\\packages\\Zygote\\rv6db\\src\\compiler\\interface2.jl:0",
      " [11] Pullback",
      "    @ ~\\Documents\\CodeSpace\\SimpleGNN\\src\\utils.jl:52 [inlined]",
      " [12] (::typeof(∂(λ)))(Δ::Float64)",
      "    @ Zygote ~\\.julia\\packages\\Zygote\\rv6db\\src\\compiler\\interface2.jl:0",
      " [13] (::Zygote.var\"#84#85\"{Zygote.Params, typeof(∂(λ)), Zygote.Context})(Δ::Float64)",
      "    @ Zygote ~\\.julia\\packages\\Zygote\\rv6db\\src\\compiler\\interface.jl:343",
      " [14] gradient(f::Function, args::Zygote.Params)",
      "    @ Zygote ~\\.julia\\packages\\Zygote\\rv6db\\src\\compiler\\interface.jl:76",
      " [15] train!(gnn::GNN², data::NamedTuple{(:node_features, :node_labels, :adjacency_list, :train_indices, :val_indices, :test_indices, :num_classes, :num_nodes, :num_edges, :directed), Tuple{Matrix{Float32}, Vector{Int64}, Vector{Vector{Int64}}, UnitRange{Int64}, UnitRange{Int64}, UnitRange{Int64}, Int64, Int64, Int64, Bool}}; epochs::Int64, patience::Int64, lr::Float64)",
      "    @ Main ~\\Documents\\CodeSpace\\SimpleGNN\\src\\utils.jl:51",
      " [16] train!(gnn::GNN², data::NamedTuple{(:node_features, :node_labels, :adjacency_list, :train_indices, :val_indices, :test_indices, :num_classes, :num_nodes, :num_edges, :directed), Tuple{Matrix{Float32}, Vector{Int64}, Vector{Vector{Int64}}, UnitRange{Int64}, UnitRange{Int64}, UnitRange{Int64}, Int64, Int64, Int64, Bool}})",
      "    @ Main ~\\Documents\\CodeSpace\\SimpleGNN\\src\\utils.jl:34",
      " [17] top-level scope",
      "    @ In[13]:3",
      " [18] eval",
      "    @ .\\boot.jl:360 [inlined]",
      " [19] include_string(mapexpr::typeof(REPL.softscope), mod::Module, code::String, filename::String)",
      "    @ Base .\\loading.jl:1094"
     ]
    }
   ],
   "source": [
    "gat = GNN²(1433,16,7;layer=GATConv,dropout=.6)\n",
    "\n",
    "train!(gat,data)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia",
   "language": "julia",
   "name": "julia-1.6"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.6.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
